from sqlalchemy import create_engine, Column, Integer, String, Date, DateTime, func
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from datetime import datetime
import logging
from typing import List, Dict, Tuple, Optional

# 创建基类
Base = declarative_base()

class LotteryDraw(Base):
    __tablename__ = 'lottery_draws'

    id = Column(Integer, primary_key=True)
    draw_num = Column(String(20), unique=True, nullable=False)
    draw_date = Column(Date, nullable=False)
    red_ball_1 = Column(Integer, nullable=False)
    red_ball_2 = Column(Integer, nullable=False)
    red_ball_3 = Column(Integer, nullable=False)
    red_ball_4 = Column(Integer, nullable=False)
    red_ball_5 = Column(Integer, nullable=False)
    red_ball_6 = Column(Integer, nullable=False)
    blue_ball = Column(Integer, nullable=False)
    created_at = Column(DateTime, default=func.now())
    updated_at = Column(DateTime, default=func.now(), onupdate=func.now())

class DLTDraw(Base):
    __tablename__ = 'dlt_draws'

    id = Column(Integer, primary_key=True)
    draw_num = Column(String(20), unique=True, nullable=False)
    draw_date = Column(Date, nullable=False)
    red_ball_1 = Column(Integer, nullable=False)
    red_ball_2 = Column(Integer, nullable=False)
    red_ball_3 = Column(Integer, nullable=False)
    red_ball_4 = Column(Integer, nullable=False)
    red_ball_5 = Column(Integer, nullable=False)
    blue_ball_1 = Column(Integer, nullable=False)
    blue_ball_2 = Column(Integer, nullable=False)
    created_at = Column(DateTime, default=func.now())
    updated_at = Column(DateTime, default=func.now(), onupdate=func.now())

class Database:
    def __init__(self, db_url: str = 'sqlite:///lottery.db'):
        self.engine = create_engine(db_url)
        Base.metadata.create_all(self.engine)
        Session = sessionmaker(bind=self.engine)
        self.session = Session()

    def _process_ssq_data(self, draw_data: Dict) -> Tuple[List[int], int]:
        """处理双色球数据
        
        Args:
            draw_data: 开奖数据字典
        Returns:
            Tuple[List[int], int]: 红球列表和蓝球
        """
        try:
            red_numbers = [int(num) for num in draw_data.get('red_numbers', [])]
            blue_number = int(draw_data.get('blue_number', 0))
            
            if len(red_numbers) != 6:
                raise ValueError(f"双色球需要6个红球，当前有{len(red_numbers)}个")
            if not blue_number:
                raise ValueError("缺少蓝球数据")
                
            return red_numbers, blue_number
        except (ValueError, TypeError) as e:
            raise ValueError(f"双色球数据处理错误: {str(e)}")

    def _process_dlt_data(self, draw_data: Dict) -> Tuple[List[int], List[int]]:
        """处理大乐透数据
        
        Args:
            draw_data: 开奖数据字典
        Returns:
            Tuple[List[int], List[int]]: 红球列表和蓝球列表
        """
        try:
            red_numbers = [int(num) for num in draw_data.get('red_numbers', [])]
            blue_numbers = [int(num) for num in draw_data.get('blue_numbers', [])]
            
            if len(red_numbers) != 5:
                raise ValueError(f"大乐透需要5个红球，当前有{len(red_numbers)}个")
            if len(blue_numbers) != 2:
                raise ValueError(f"大乐透需要2个蓝球，当前有{len(blue_numbers)}个")
                
            return red_numbers, blue_numbers
        except (ValueError, TypeError) as e:
            raise ValueError(f"大乐透数据处理错误: {str(e)}")

    def save_draw(self, draw_data: Dict, lottery_type: str = 'ssq') -> bool:
        """保存单条开奖记录
        
        Args:
            draw_data: 开奖数据字典
            lottery_type: 彩票类型，'ssq'为双色球，'dlt'为大乐透
        """
        try:
            # 验证基本数据
            if not draw_data.get('draw_num'):
                raise ValueError("缺少开奖期号")
            if not draw_data.get('draw_date'):
                raise ValueError("缺少开奖日期")

            # 转换日期字符串为日期对象
            draw_date = datetime.strptime(draw_data['draw_date'], '%Y-%m-%d').date()
            
            # 根据彩票类型选择对应的模型类和处理方法
            if lottery_type == 'ssq':
                model_class = LotteryDraw
                red_numbers, blue_number = self._process_ssq_data(draw_data)
            else:
                model_class = DLTDraw
                red_numbers, blue_numbers = self._process_dlt_data(draw_data)
            
            # 检查记录是否已存在
            existing_draw = self.session.query(model_class).filter_by(
                draw_num=draw_data['draw_num']
            ).first()

            if existing_draw:
                # 更新现有记录
                existing_draw.draw_date = draw_date
                for i, ball in enumerate(red_numbers, 1):
                    setattr(existing_draw, f'red_ball_{i}', ball)
                
                if lottery_type == 'ssq':
                    existing_draw.blue_ball = blue_number
                else:
                    existing_draw.blue_ball_1 = blue_numbers[0]
                    existing_draw.blue_ball_2 = blue_numbers[1]
            else:
                # 创建新记录
                draw_dict = {
                    'draw_num': draw_data['draw_num'],
                    'draw_date': draw_date
                }
                
                # 添加红球
                for i, ball in enumerate(red_numbers, 1):
                    draw_dict[f'red_ball_{i}'] = ball
                
                # 添加蓝球
                if lottery_type == 'ssq':
                    draw_dict['blue_ball'] = blue_number
                else:
                    draw_dict['blue_ball_1'] = blue_numbers[0]
                    draw_dict['blue_ball_2'] = blue_numbers[1]
                
                new_draw = model_class(**draw_dict)
                self.session.add(new_draw)

            self.session.commit()
            return True

        except Exception as e:
            self.session.rollback()
            logging.error(f"保存开奖记录失败: {str(e)}")
            return False

    def save_draws(self, draws_data: List[Dict], lottery_type: str = 'ssq') -> int:
        """批量保存开奖记录
        
        Args:
            draws_data: 开奖数据列表
            lottery_type: 彩票类型，'ssq'为双色球，'dlt'为大乐透
        """
        success_count = 0
        for draw_data in draws_data:
            if self.save_draw(draw_data, lottery_type):
                success_count += 1
        return success_count

    def get_latest_draw(self, lottery_type: str = 'ssq') -> Optional[LotteryDraw]:
        """获取最新一期开奖记录
        
        Args:
            lottery_type: 彩票类型，'ssq'为双色球，'dlt'为大乐透
        Returns:
            Optional[LotteryDraw]: 最新一期开奖记录，如果没有则返回None
        """
        model_class = LotteryDraw if lottery_type == 'ssq' else DLTDraw
        return self.session.query(model_class).order_by(
            model_class.draw_date.desc()
        ).first()

    def close(self):
        """关闭数据库连接"""
        self.session.close()